// This file is licensed under the Elastic License 2.0. Copyright 2021-present, StarRocks Limited.

package com.starrocks.sql.optimizer.base;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.starrocks.thrift.TDistributionType;

import java.util.List;
import java.util.Map;

public class DistributionSpec {
    protected final DistributionType type;
    protected PropertyInfo propertyInfo;

    protected DistributionSpec(DistributionType type) {
        this(type, new PropertyInfo());
    }

    protected DistributionSpec(DistributionType type, PropertyInfo propertyInfo) {
        this.type = type;
        this.propertyInfo = propertyInfo;
    }

    // Property information for hash distribution desc, it used for check DistributionSpec satisfy condition.
    public static final class PropertyInfo {
        public long tableId = -1;
        public boolean isReplicate = false;
        public List<Long> partitionIds = Lists.newArrayList();
        // record nullable columns generated by outer join
        public ColumnRefSet nullableColumns = new ColumnRefSet();
        // col -> all equivalent join columns
        // For example, if A.a, B.b and C.c are equivalent, then we have three key value paris
        // A.a -> [A.a, B.b, C.c]
        // B.b -> [A.a, B.b, C.c]
        // C.c -> [A.a, B.b, C.c]
        private final Map<Integer, ColumnRefSet> joinEquivalentColumns = Maps.newHashMap();

        public boolean isSinglePartition() {
            return partitionIds.size() == 1;
        }

        public boolean isEmptyPartition() {
            return partitionIds.size() == 0;
        }

        public void addJoinEquivalentPair(int leftColumn, int rightColumn) {
            ColumnRefSet leftColumns = null;
            ColumnRefSet rightColumns = null;
            leftColumns = joinEquivalentColumns.get(leftColumn);
            rightColumns = joinEquivalentColumns.get(rightColumn);
            ColumnRefSet columns;
            if (leftColumns == null && rightColumns == null) {
                columns = new ColumnRefSet();
                joinEquivalentColumns.put(leftColumn, columns);
                joinEquivalentColumns.put(rightColumn, columns);
            } else if (leftColumns == null) {
                columns = rightColumns;
                joinEquivalentColumns.put(leftColumn, columns);
            } else if (rightColumns == null) {
                columns = leftColumns;
                joinEquivalentColumns.put(rightColumn, columns);
            } else {
                leftColumns.union(rightColumns);
                columns = leftColumns;
                joinEquivalentColumns.put(rightColumn, columns);
            }
            columns.union(leftColumn);
            columns.union(rightColumn);
        }

        public boolean isEquivalentJoinOnColumns(int leftColumn, int rightColumn) {
            if (joinEquivalentColumns.containsKey(leftColumn)) {
                return joinEquivalentColumns.get(leftColumn).contains(rightColumn);
            }
            return false;
        }

        public int[] getEquivalentJoinOnColumns(int column) {
            if (joinEquivalentColumns.containsKey(column)) {
                return joinEquivalentColumns.get(column).getColumnIds();
            } else {
                return new int[0];
            }
        }
    }

    @SuppressWarnings("unchecked")
    public <T> T cast() {
        return (T) this;
    }

    public PropertyInfo getPropertyInfo() {
        return propertyInfo;
    }

    public DistributionType getType() {
        return type;
    }

    public static DistributionSpec createAnyDistributionSpec() {
        return new AnyDistributionSpec();
    }

    public static HashDistributionSpec createHashDistributionSpec(HashDistributionDesc distributionDesc) {
        return new HashDistributionSpec(distributionDesc);
    }

    public static DistributionSpec createReplicatedDistributionSpec() {
        return new ReplicatedDistributionSpec();
    }

    public static DistributionSpec createGatherDistributionSpec() {
        return new GatherDistributionSpec();
    }

    public static DistributionSpec createGatherDistributionSpec(long limit) {
        return new GatherDistributionSpec(limit);
    }

    public boolean isSatisfy(DistributionSpec spec) {
        return false;
    }

    public enum DistributionType {
        ANY,
        BROADCAST,
        SHUFFLE,
        GATHER,
        ;

        public TDistributionType toThrift() {
            if (this == ANY) {
                return TDistributionType.ANY;
            } else if (this == BROADCAST) {
                return TDistributionType.BROADCAST;
            } else if (this == SHUFFLE) {
                return TDistributionType.SHUFFLE;
            } else {
                return TDistributionType.GATHER;
            }
        }
    }

    @Override
    public String toString() {
        return type.toString();
    }
}
